 
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.indices.common.struct_store.sql import SQLStructDatapointExtractor
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer, \
    DocumentSummaryIndex
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel
from llama_index.core.indices.property_graph.base import PropertyGraphIndex
from llama_index.core.indices.property_graph.retriever import PGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.base import BasePGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.custom import (
    CustomPGRetriever,
    CUSTOM_RETRIEVE_TYPE,
)
from llama_index.core.indices.property_graph.sub_retrievers.cypher_template import (
    CypherTemplateRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.llm_synonym import (
    LLMSynonymRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.text_to_cypher import (
    TextToCypherRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.vector import (
    VectorContextRetriever,
)
from llama_index.core.indices.property_graph.transformations.implicit import (
    ImplicitPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.schema_llm import (
    SchemaLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.simple_llm import (
    SimpleLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.dynamic_llm import (
    DynamicLLMPathExtractor,
)
from llama_index.core.indices.property_graph.utils import default_parse_triplets_fn

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm

from pathlib import Path
from typing import List

from llama_index.core.indices.composability.graph import ComposableGraph
from llama_index.core.indices.list.base import SummaryIndex
from llama_index.core.indices.loading import load_graph_from_storage
from llama_index.core.indices.vector_store.base import VectorStoreIndex
from llama_index.core.schema import Document
from llama_index.core.storage.storage_context import StorageContext


storage_context = StorageContext.from_defaults()
doc001=Document(text='''　用电量是经济运行的“晴雨表”“风向标”。国家能源局数据显示，8月全社会用电量10154亿千瓦时，同比增长5.0%。这是继今年7月用电量首次突破万亿千瓦时之后，再度破万亿，在全球也属首次。

　　“1万亿千瓦时”，相当于由乌东德、白鹤滩、三峡等6座梯级水电站构成的世界最大清洁能源走廊全年发电量的3倍多，也约是我国2015年7月用电量的2倍。

　　中国电力企业联合会统计与数智部副主任蒋德斌认为，月度用电量连续两月突破万亿千瓦时的背后，有今夏高温高湿天气来得早持续久，带动居民用电快速增长的拉动，也显示出新质生产力蓬勃发展，正形成新的经济增长点，推动用电量向上攀升。''')

doc002=Document(text='''当地时间9日，美国国会参议院第七次尝试表决通过临时拨款法案，但共和党和民主党分别提出的临时拨款法案草案仍然均未获通过，尽快结束政府“停摆”的努力再次失败。

美国政府停摆进入第10天，对美国社会各方面影响已经开始显现。受影响最大的是哪些领域？美国政府关门还会持续多久？如果问题迟迟得不到解决，美国还会受到怎样的冲击？''')

 

# construct index
vector_index_1 = VectorStoreIndex.from_documents(
documents=[doc001],
storage_context=storage_context,
)

# construct second index, testing vector store overlap
vector_index_2 = VectorStoreIndex.from_documents(
documents=[doc002],
storage_context=storage_context,
)

 

# construct graph
graph = ComposableGraph.from_indices(
SummaryIndex,
children_indices=[vector_index_1, vector_index_2],
index_summaries=["关于社会用电量", "关于美国政府停摆"],
storage_context=storage_context,
)

query_engine = graph.as_query_engine()
response = query_engine.query("8月全社会用电量有多少？")
print(response)
response = query_engine.query("美国政府停摆几天？")
print(response)
 
 